import json
import os
import SimpleITK as sitk
import numpy as np
import matplotlib.pyplot as plt


def Gen4ContrastMRI(datalist, target='./'):
    """
    分别生成四个模态 MRI 影像的 json 文件
    :param data_root: 数据集根目录
    :param target: 生成文件的输出地址
    :return:
    """    
    with open(datalist) as f:
        json_data = json.load(f)

    print(json_data.keys())
    
    # for m in ['t1', 't1ce', 'flair', 't2']:
    for i in json_data['training']:
        del i['image'][0:2]
        del i['image'][1:]


    with open(os.path.join(target, 't1.json'), 'w') as f:
        json.dump(json_data, f)
        

def visual_slice(images, images_name):
    plt.imsave(images_name, images, cmap='gray', dpi=150)

def visual_hot_map(images, images_name):
    plt.imsave(images_name, images, cmap='RdBu_r', dpi=150)
    
def visual_seg(images, images_name):
    plt.imsave(images_name, images, cmap='gray', dpi=150)

def read_nii(path):
    image = sitk.ReadImage(path)
    image = sitk.GetArrayFromImage(image)
    return image     

def norm(image):
    return (image - np.min(image)) / (np.max(image) - np.min(image))   

if __name__ == '__main__':
    t1_path = '/Users/qlc/Desktop/Dataset/Brats2021/TrainingData/BraTS2021_00018/BraTS2021_00018_t1.nii.gz'
    t1ce_path = '/Users/qlc/Desktop/Dataset/Brats2021/TrainingData/BraTS2021_00018/BraTS2021_00018_t1ce.nii.gz'
    t2_path = '/Users/qlc/Desktop/Dataset/Brats2021/TrainingData/BraTS2021_00018/BraTS2021_00018_t2.nii.gz'
    flair_path = '/Users/qlc/Desktop/Dataset/Brats2021/TrainingData/BraTS2021_00018/BraTS2021_00018_flair.nii.gz'
    seg_path = '/Users/qlc/Desktop/Dataset/Brats2021/TrainingData/BraTS2021_00018/BraTS2021_00018_seg.nii.gz'
    
    t1 = read_nii(t1_path)
    t2 = read_nii(t2_path)
    flair = read_nii(flair_path)
    t1ce = read_nii(t1ce_path)
    seg = read_nii(seg_path)
    
    t1_n = norm(t1)
    t2_n = norm(t2)
    flair_n = norm(flair)
    t1ce_n = norm(t1ce)
    
    seg_1 = seg==1
    seg_2 = seg==2
    seg_4 = seg==4
        
    # # slice 86
    visual_hot_map(norm(np.abs(t1ce_n - 1) - t1_n)[86], './t1-t1ce.png')
    visual_hot_map(norm(flair - np.abs(t2 - 1))[86], './flair-t2.png')
    visual_hot_map(norm(t1 - np.abs(t2 - 1))[86], './t1-t2.png')
    
    visual_slice(norm(t1_n - np.abs(t1ce_n - 1))[86] > 0.55, './t1-t1ce-0.55.png')
    visual_slice(norm(flair - np.abs(t2 - 1))[86] > 0.75, './flair-t2-0.55.png')
    visual_slice(norm(t1 - np.abs(t2 - 1))[86] > 0.5, './t1-t2-0.55.png')
    
    print(np.max(norm(t1_n - np.abs(t1ce_n - 1))[86]), np.min(norm(t1_n - np.abs(t1ce_n - 1))[86]))
    
    # # slice 100
    # visual_hot_map(norm(t1_n - np.abs(t1ce_n - 1))[100], './t1-t1ce.png')
    # visual_hot_map(norm(flair - np.abs(t2 - 1))[100], './flair-t2.png')
    # visual_hot_map(norm(t1 - np.abs(t2 - 1))[100], './t1-t2.png')
    
    # visual_slice(norm(t1_n - np.abs(t1ce_n - 1))[100] > 0.55, './t1-t1ce-0.55.png')
    # visual_slice(norm(flair - np.abs(t2 - 1))[100] > 0.75, './flair-t2-0.55.png')
    # visual_slice(norm(t1 - np.abs(t2 - 1))[100] > 0.5, './t1-t2-0.55.png')

    
    # visual_slice(np.power(norm(np.abs(t1ce_n - 1) - t1_n)[100], 2) > 0.99, './b.png')
    
    # visual_hot_map(norm(flair_n - t2_n)[100], './b.png')
    
    
    
    
    
    
    
    
    
    
    
    # j = '/Users/qlc/Code/Gmim/BraTS21/jsons/brats21_folds.json'
    # Gen4ContrastMRI(j)
    
    # seg = '/Users/qlc/Desktop/Dataset/Brats2023/Adult_Glioma/TrainingData/BraTS-GLI-00000-000/BraTS-GLI-00000-000-seg.nii.gz'
    
    # seg = sitk.ReadImage(seg)
    # seg = sitk.GetArrayFromImage(seg)
    
    # # print(np.where(seg == 1))
    # # print(np.where(seg==2))
    # # print(np.where(seg==3))
    # # print(seg[84, 88, 142])
    # # where1 = np.where(seg==1)

    # # print(np.unique(where1[0]))
    # # print(np.unique(where1[1]))
    # # print(np.unique(where1[2]))
    
    # import torch
    
    # seg_t = torch.from_numpy(seg)
    # loc  = torch.where(seg_t==1)
    
    # print(loc[0])
    # print(loc[1])
    # print(loc[2])
    
    # print(loc[0].size())
    # print(loc[1].size())
    # print(loc[2].size())
    
    # print(loc[0][11737])
    # # print(seg_t[loc[0][11737], loc[1][0], loc[2][0]])
    # # print(torch.where(seg_t==1))
    # # print()
    

    